import FileEnum   as fe
import GeotiffMan as gm
import os.path
import numpy

#
# Invalid data in input data frames are still unknows yet,
# so, consider it as zero, all calculations have codes to
# avoid problems caused by zero-divider.
#
# If invalid data in input data frames are other values,
# such codes should be modified to operate properly.
#


#--------------------------------------------------------------
#
# For VHI Drought (1.1.22)
#


def calc_ndvi(nir, red, invalid_value=0):
    '''
    NDVI - Normalized Difference Vegetation Index，植被指数
    
           NIR - RED
    NDVI = ---------
           NIR + RED

    NIR - Near Infrared
    red - Red
    
    Above is for each pixel.
    
    Input:
        nir - NIR data frame, in type of numpy array
        red - RED data frame, in type of numpy array
        
    Output:
        ndvi - NDVI data frame, type of numpy ndarray
    '''
    sub = nir - red
    add = nir + red
    #
    # For invalid pixels, both nir and red are zeros,
    # thus nir + red are also zeros. 
    # Calculate NDVI of these pixels will cause divide zero errors, 
    # and get the results of nans(Not A Number).
    # To avoid above issues, add 1 to the nir + red for invalid pexels
    # before use it as divider.
    #
    invalid = add == invalid_value
    ndvi = sub / (add + invalid)
    valid = add != invalid_value
    return ndvi * valid # now invalid value is zero.

'''
ndvi_files = fe.enum_files('ndvi'x = numpy.array([[0, 5, 2, 3],
                 [0, 4, 5, 6],
                 [0, 7, 8, 9]]))
ndvi_files = [os.path.realpath(file) for file in ndvi_files]
ndvi_files = [file.replace('\\', '/') for file in ndvi_files]
print(ndvi_files)

nir = gm.read_geotiff_band(ndvi_files[0])
ir = gm.read_geotiff_band(ndvi_files[1])

x = calc_ndvi(nir, ir)


x = numpy.array([[0,1,2,3],
               [0,4,5,6],
               [0,7,8,9]])

y = numpy.array([[0,4,4,4],
               [0,6,6,6],
               [0,5,5,5]])

xpy = x + y
xsy = x - y

zxpy = xpy == 0

n = xsy / (zxpy + xpy)
'''

def get_ndvi_min(ndvi_list):
    '''
    Get minimum ndvi value of each pixel.
    
    Input :
        NDVI data frames, in type of numpy array
        
    Output :
        Minimum NDVI data frame, in type of numpy array
    '''
    if len(ndvi_list) == 0:
        print('Zero length input list of get_ndvi_min()')
        0/0
    
    result = numpy.ones(ndvi_list[0].shape, ndvi_list[0].dtype)
    
    for ndvi in ndvi_list:
        i_ndvi   = ndvi < result
        i_result = i_ndvi == False
        result = ndvi * i_ndvi + result * i_result
    
    return result

    
def get_ndvi_max(ndvi_list):
    '''
    Get maximum ndvi value of each pixel.
    
    Input :
        NDVI data frames, in type of numpy array
        
    Output :
        maximum NDVI data frame, in type of numpy array
    '''
    if len(ndvi_list) == 0:
        print('Zero length input list of get_ndvi_min()')
        0/0
    
    result = numpy.zeros(ndvi_list[0].shape, ndvi_list[0].dtype)
    
    for ndvi in ndvi_list:
        i_ndvi   = ndvi > result
        i_result = i_ndvi == False
        result = ndvi * i_ndvi + result * i_result
    
    return result


    
def calc_vci(ndvi_list):
    '''
    VCI - Vegetation Condition Index, 植被状态指数
    
          NDVI     - NDVI_min
    VCI = -------------------
          NDVI_max - NDVI_min
    
    
    '''
    ndvi_min = get_ndvi_min(ndvi_list)
    #print('ndvi_min\n', ndvi_min)
    
    ndvi_max = get_ndvi_max(ndvi_list)
    #print('ndvi_max\n', ndvi_max)
    
    vci_list = []
    
    for ndvi in ndvi_list:
        sub_up   = ndvi     - ndvi_min
        sub_down = ndvi_max - ndvi_min
        #
        # For invalid pixels, both nir and red are zeros,
        # thus nir + red are also zeros. 
        # Calculate NDVI of these pixels will cause divide zero errors, 
        # and get the results of nans(Not A Number).
        # To avoid above issues, add 1 to the nir + red for invalid pexels
        # before use it as divider.
        #
        zeros = sub_down == 0
        vci_list.append(sub_up / (sub_down + zeros))
    
    #print('vci_list\n', vci_list)
    return vci_list


'''
x = numpy.array([[0, 5, 2, 3],
                 [0, 4, 5, 6],
                 [0, 7, 8, 9]])

y = numpy.array([[0, 4, 4, 4],
                 [0, 6, 6, 6],
                 [0, 5, 5, 5]])
    
z = numpy.array([[0, 8, 4, 2],
                 [0, 3, 6, 4],
                 [0, 7, 3, 2]])

z = calc_vci([x/10, y/10, z/10])
#'''

#--------------------------------------------------------------

def calc_brightness_temperature(radiation, k1, k2):
    '''
    Calculate brightness temperature(亮温) from radiation(辐亮度)
    
                                    K2
    brightness_temperature = ----------------
                             ln(K1/lamda + 1)
         
    K1 K2 - constants,定标常数
    lamda - 热红外波段的辐亮度
    
    Input:
        radiation - 热红外波段的辐亮度, data frame, in type of numpy array
        k1 - constant1
        k2 - constant2
    
    Output:
        brightness_temperature - data frame, in type of numpy array
    '''
    zeros = radiation == 0
    valid = radiation != 0
    
    radiation += zeros
    bt = k2 / numpy.log2(k1/radiation + 1) * valid
    
    return bt

'''
x = numpy.array([[0, 5, 2, 3],
                 [0, 4, 5, 6],
                 [0, 7, 8, 9]])
    
z = calc_brightness_temperature(x, 1, 1)
'''


def get_bt_min(bt_list):
    '''
    See get_ndvi_min.
    '''
    return get_ndvi_min(bt_list)


def get_bt_max(bt_list):
    '''
    See get_ndvi_max.
    '''    
    return get_ndvi_max(bt_list)


def calc_tci(bt_list):
    '''
    Calculate Temperature Condition Index, TCI,温度状况指数
    
          BT_max - BT
    TCI = ---------------
          BT_max - BT_min
          
    Input:
        bt_list - brightness temperature data frames,
                  each frame is in type of numpy array
    Output:
        temperature condition index
    '''
    bt_min = get_bt_min(bt_list)
    print('bt_min\n', bt_min)
    
    bt_max = get_bt_max(bt_list)
    print('bt_max\n', bt_max)
    
    tci_list = []
    
    for bt in bt_list:
        sub_up   = bt_max - bt
        sub_down = bt_max - bt_min
        #
        # For invalid pixels, both nir and red are zeros,
        # thus nir + red are also zeros. 
        # Calculate NDVI of these pixels will cause divide zero errors, 
        # and get the results of nans(Not A Number).
        # To avoid above issues, add 1 to the nir + red for invalid pexels
        # before use it as divider.
        #
        zeros = sub_down == 0
        tci_list.append(sub_up / (sub_down + zeros))
    
    #print('vci_list\n', vci_list)
    return tci_list
    
'''
x = numpy.array([[0, 5, 2, 3],
                 [0, 4, 5, 6],
                 [0, 7, 8, 9]])

y = numpy.array([[0, 4, 4, 4],
                 [0, 6, 6, 6],
                 [0, 5, 5, 5]])
    
z = numpy.array([[0, 8, 4, 2],
                 [0, 3, 6, 4],
                 [0, 7, 3, 2]])

z = calc_tci([x/10, y/10, z/10])
#'''
    

def calc_vhi(vci, k_vci, tci, k_tci):
    '''
    Calculate vegetation healthy index, VHI, 植被健康度指数
    
    VHI = vci * k_vci + tci * k_tci
    
    Input:
        vci   - VCI data frame
        k_vci - coef. of vci
        tci   - TCI data frame
        k_tci - coef. of tci
    Output:
        vegetation healthy index
    '''
    return vci * k_vci + tci * k_tci
    

#----------------------------------------------------------
# For NDVI Drought (1.1.24)
#    

def get_ndvi_mean(ndvi_list):
    '''
    Calc mean ndvi of multi ndvi data frames.
    
    Input:
        ndvi_list - list of ndvi data frames.
        ndvi data frame is in type of numpy array
    Output:
        mean ndvi data frame
    '''
    result = numpy.zeros(ndvi_list[0].shape, ndvi_list[0].dtype)
    
    for ndvi in ndvi_list:
        result += ndvi
    
    return result / len(ndvi_list)


def calc_ndvi_drought(ndvi_list, grades=4):
    '''
    Calculate NDVI drought classification.
    
    每个像元逐期数据减去其多年平均值。
    按照逐年最小值作为下限，0为上限，等间隔分为4段；
    每期每像元大于0者定为等级0（无干旱）；
    小于0者，在上述4段中哪个段对应分级就可。

    Input:
        dvi_list - list of ndvi data frames.
        ndvi data frame is in type of numpy array
    Output:
        classificaiton result
    '''
    ndvi_mean = get_ndvi_mean(ndvi_list)
    ndvi_min  = get_ndvi_min(ndvi_list)
    ndvi_low_range = ndvi_min - ndvi_mean
    ndvi_seg_range = ndvi_low_range / grades
    
    classified = []
    
    for ndvi in ndvi_list:
        ndvi_centered = ndvi - ndvi_mean
        ndvi_centered = ndvi_centered * (ndvi_centered < 0)
        classx = ndvi_centered / ndvi_seg_range
        classified.append(numpy.ceil(classx))
    
    return classified
        
'''
x = numpy.array([[0, 5, 2, 3],
                 [0, 4, 5, 6],
                 [0, 7, 8, 9]])

y = numpy.array([[0, 4, 4, 4],
                 [0, 6, 6, 6],
                 [0, 5, 5, 5]])
    
z = numpy.array([[0, 8, 4, 2],
                 [0, 3, 6, 4],
                 [0, 7, 3, 2]])

z = calc_tci([x/10, y/10, z/10])

x = z/3
x = numpy.ceil(x)
print(x)
#'''    
    